import base64
import json
import shutil
import uuid
import hashlib

from core import get_info_handler, set_info_handler
from lib.Event import Event
from lib.lib import *
import os

from lib.support_lib import run_command_with_return, CustomJsonEncoder

CHUNK_SIZE = 1000 * 1024  # 1MB


Info = {
    "name": "外部指令处理器",
    "id": "external_handler",
    "dependencies": [],
    "pip_dependencies": [],
    "linux_dependencies": []
}


def try_decode_bytes(data: bytes, encodings=('utf-8', 'gbk')) -> tuple[str, str]:
    for enc in encodings:
        try:
            return enc, data.decode(enc)
        except Exception:
            continue
    return 'gbk', data.decode('gbk', errors='ignore')


def calculate_sha256(data: bytes) -> str:
    if not isinstance(data, bytes):
        raise TypeError("Input for calculate_sha256 must be bytes.")
    return hashlib.sha256(data).hexdigest()


def calculate_file_sha256(file_path: str) -> str:
    sha256_hash = hashlib.sha256()
    try:
        with open(file_path, "rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        return sha256_hash.hexdigest()
    except FileNotFoundError:
        return ""


def b64_encode_bytes(b: bytes) -> str:
    return base64.b64encode(b).decode()


def b64_decode(txt, file):
    data = base64.b64decode(txt)
    with open(file, 'wb') as output:
        output.write(data)
    return f"文件 {file} 写入成功\n"


def get_running():
    running = {}
    for module in static["running"].keys():
        try:
            running[loaded_plugins[module].Info["name"]] = static["running"][module]
        except KeyError:
            running[module] = static["running"][module]
    return running


class ExternalHandlerLogic:
    def __init__(self):
        self.upload_sessions = {}
        self.download_sessions = {}
        static["event_handler"].register_event("ExternalIO_IN", self.process_input_event)

    def process_input_event(self, event: Event):
        if event.name != "ExternalIO_IN":
            return

        command_message = event.data.get("message")
        client_id = event.data.get("client_id")

        if not command_message:
            error_message = "错误: 接收到空消息\n"
            static["event_handler"].trigger_event(
                Event("ExternalIO_ERROR", error_message=error_message,
                      original_command="", client_id=client_id)
            )
            return

        result = self._handle_command_and_return(command_message, client_id)
        if result is not None:
            static["event_handler"].trigger_event(
                Event("ExternalIO_OUT", response_data=result, client_id=client_id)
            )

    def _handle_command_and_return(self, command_message: str, client_id: str):
        try:
            ctrl = ord(command_message[0])
            message = command_message[1:]

            if ctrl == 0x01:
                return chr(0x01) + run_command_with_return(message)

            elif ctrl == 0x02:
                uid = message[0:4]
                message = message[4:]
                message_parts = message.split(' ', 1)
                if len(message_parts) < 1:
                    response_content = "参数不足"
                else:
                    sub_command = message_parts[0]
                    sub_args = message_parts[1].split(' ') if len(message_parts) > 1 else []
                    if sub_command == "get":
                        response_content = get_info_handler(sub_args)
                    elif sub_command == "set":
                        response_content = set_info_handler(sub_args)
                    else:
                        response_content = "unknown"
                return chr(0x02) + uid + response_content

            elif ctrl == 0x03:
                cuid = message[0:4]
                data = {
                    "/T/V系统": static["SYS_INFO"],
                    "/APython版本: ": static["PY_VER"],
                    "/P在线模式": static["online"],
                    "/P/AGit可用": static["git_available"],
                    "/L/A访问项目": "https://github.com/furryaxw/AFEDIUM/",
                    "/C/Vmodules": get_running(),
                }
                return chr(0x03) + cuid + json.dumps(data, indent=2, cls=CustomJsonEncoder)

            elif ctrl == 0x11:  # Upload Handshake / Small File Upload
                parts = message.split(' ', 3)

                # Case 1: Small file upload with checksum and encoding (from file_detail_screen)
                # Format: file_path base64_data expected_md5_checksum encoding
                if len(parts) == 4:
                    file_path, base64_data, expected_checksum, encoding = parts
                    try:
                        data_bytes = base64.b64decode(base64_data)
                        calculated_checksum = calculate_sha256(data_bytes)
                        if calculated_checksum != expected_checksum:
                            return chr(0x11) + f"ERROR_CHECKSUM_MISMATCH: {file_path} 校验失败"

                        with open(file_path, 'wb') as f:
                            f.write(data_bytes)
                        return chr(0x11) + f"{file_path} 保存成功"
                    except Exception as e:
                        return chr(0x11) + f"保存失败: {e}"
                # Case 2: Large file upload handshake (from file_manager_screen.dart)
                # Format: file_path total_chunks expected_sha256
                elif len(parts) == 3 and parts[1].isdigit():
                    file_path, chunk_count_str, expected_sha256 = parts
                    try:
                        chunk_count = int(chunk_count_str)
                    except ValueError:
                        return chr(0x11) + "参数错误: chunk_count 无效"

                    session_id = str(uuid.uuid4())
                    self.upload_sessions[session_id] = {
                        "file_path": file_path,
                        "chunks": chunk_count,
                        "received_chunk_indices": set(),
                        "received_chunks_data": {},
                        "expected_sha256": expected_sha256,
                        "client_id": client_id
                    }
                    return chr(0x11) + session_id
                # Case 3: Legacy small file upload (if still supported, without checksum/encoding)
                # Format: file_path base64_data
                elif len(parts) == 2:
                    file_path = parts[0]
                    base64_data = parts[1]
                    try:
                        data = base64.b64decode(base64_data)
                        with open(file_path, 'wb') as f:
                            f.write(data)
                        return chr(0x11) + f"{file_path} 保存成功"
                    except Exception as e:
                        return chr(0x11) + f"保存失败: {e}"
                else:
                    return chr(0x11) + "参数错误"

            elif ctrl == 0x18:  # Chunk Transfer (Upload & Download)
                parts = message.split(' ', 2)
                if len(parts) == 3:  # Uploading a chunk (session_id, chunk_index, base64_chunk)
                    session_id, chunk_index_str, base64_chunk = parts
                    session = self.upload_sessions.get(session_id)
                    if not session:
                        return chr(0x18) + "无效上传会话"

                    try:
                        chunk_index = int(chunk_index_str)
                    except ValueError:
                        return chr(0x18) + "参数错误: chunk_index 无效"

                    session["received_chunk_indices"].add(chunk_index)
                    session["received_chunks_data"][chunk_index] = base64.b64decode(base64_chunk)

                    if len(session["received_chunk_indices"]) == session["chunks"]:
                        full_data_bytes = b''
                        for i in range(session["chunks"]):
                            if i not in session["received_chunks_data"]:
                                print(f"[WARNING] Missing chunk {i} for session {session_id} during assembly.")
                                return chr(0x18) + f"文件组装失败: 缺失分块"
                            full_data_bytes += session["received_chunks_data"][i]

                        file_path = session["file_path"]
                        with open(file_path, 'wb') as f:
                            f.write(full_data_bytes)

                        calculated_sha256 = calculate_sha256(full_data_bytes)
                        expected_sha256 = session["expected_sha256"]

                        del self.upload_sessions[session_id]

                        if calculated_sha256 == expected_sha256:
                            response_message = chr(0x18) + f"END {calculated_sha256}"
                        else:
                            response_message = chr(0x18) + f"ERROR_SHA_MISMATCH {calculated_sha256}"

                        static["event_handler"].trigger_event(
                            Event("ExternalIO_OUT", response_data=response_message, client_id=client_id)
                        )
                        return None
                    return None
                else:
                    return chr(0x18) + "参数错误"

            elif ctrl == 0x19:  # Missing Chunk Re-negotiation (for Download only, re-sends chunks)
                parts = message.split(' ', 1)
                if len(parts) == 2:
                    session_id, missing_indices_str = parts
                    try:
                        missing_indices = json.loads(missing_indices_str)
                    except json.JSONDecodeError:
                        return chr(0x19) + "参数错误: missing_indices 无效JSON"

                    download_session = self.download_sessions.get(session_id)
                    if download_session:
                        file_path = download_session["file"]
                        for chunk_index in missing_indices:
                            if session_id not in self.download_sessions:
                                print(
                                    f"[ExternalHandler] Download session {session_id} cancelled during re-send.")
                                break
                            with open(file_path, 'rb') as f:
                                f.seek(chunk_index * CHUNK_SIZE)
                                chunk = f.read(CHUNK_SIZE)
                                if chunk:
                                    encoded_chunk = base64.b64encode(chunk).decode()
                                    static["event_handler"].trigger_event(
                                        Event("ExternalIO_OUT", response_data=chr(
                                            0x18) + f"{session_id} {chunk_index} {encoded_chunk}",
                                              client_id=client_id)
                                    )
                        if session_id in self.download_sessions:
                            static["event_handler"].trigger_event(
                                Event("ExternalIO_OUT", response_data=chr(0x18) + "END", client_id=client_id)
                            )
                        return None
                return chr(0x19) + "参数错误"

            elif ctrl == 0x1A:  # New: Cancel Download
                session_id = message.strip()
                if session_id in self.download_sessions:
                    print(f"[ExternalHandler] Client {client_id} cancelling download session: {session_id}")
                    del self.download_sessions[session_id]
                    static["event_handler"].trigger_event(
                        Event("ExternalIO_OUT", response_data=chr(0x1A) + "CANCELLED", client_id=client_id)
                    )
                return None

            elif ctrl == 0x1B:  # New: Start Large File Transfer (Client confirms download)
                session_id = message.strip()
                download_session = self.download_sessions.get(session_id)
                if download_session:
                    file_path = download_session["file"]
                    chunk_count = download_session["chunks"]
                    print(
                        f"[ExternalHandler] Client {client_id} confirmed download for session {session_id}. Starting chunk transfer.")
                    self._send_file_chunks(session_id, file_path, chunk_count, client_id)
                else:
                    print(f"[ExternalHandler] Invalid session ID {session_id} for 0x1B command.")
                    return chr(0x1B) + "ERROR: Invalid session ID"
                return None

            elif ctrl == 0x13:
                file_path = message.strip()
                if not os.path.exists(file_path):
                    return chr(0x13) + f"错误: 文件不存在 {file_path}"
                else:
                    file_size = os.path.getsize(file_path)

                    if file_size > CHUNK_SIZE:
                        file_sha256 = calculate_file_sha256(file_path)
                        chunk_count = (file_size + CHUNK_SIZE - 1) // CHUNK_SIZE
                        session_id = str(uuid.uuid4())
                        self.download_sessions[session_id] = {
                            "file": file_path, "chunks": chunk_count, "file_sha256": file_sha256,
                            "client_id": client_id
                        }
                        return chr(0x13) + f"{session_id}:{chunk_count}:{file_sha256}"

                    else:
                        try:
                            with open(file_path, 'rb') as f:
                                content = f.read()

                            encoding_type = 'binary'
                            if file_path.lower().endswith(('.txt', '.md', '.json', '.csv', '.log', '.xml')):
                                try:
                                    encoding, decoded_text = try_decode_bytes(content)
                                    encoding_type = encoding
                                    encoded_b64 = base64.b64encode(decoded_text.encode(encoding)).decode()
                                except Exception:
                                    encoded_b64 = base64.b64encode(content).decode()
                            else:
                                encoded_b64 = base64.b64encode(content).decode()

                            file_sha256 = calculate_sha256(content)
                            return chr(0x13) + f"{encoding_type}:{encoded_b64}:{file_sha256}"
                        except Exception as e:
                            err_msg = f"Error reading file: {e}".encode()
                            return chr(0x13) + f"utf-8:{base64.b64encode(err_msg).decode()}:"

            elif ctrl == 0x12:
                target = message.strip()
                if os.path.isdir(target):
                    shutil.rmtree(target)
                    return chr(0x12) + f"目录已删除: {target}"
                elif os.path.isfile(target):
                    os.remove(target)
                    return chr(0x12) + f"文件已删除: {target}"
                else:
                    return chr(0x12) + f"路径不存在: {target}"

            elif ctrl == 0x14:
                path_ls = message.strip() or './'
                if not os.path.isdir(path_ls):
                    return chr(0x14) + json.dumps({"error": f"{path_ls} 不是有效目录"})
                else:
                    files = os.listdir(path_ls)
                    ans = {}
                    for file in files:
                        full_path = os.path.join(path_ls, file)
                        permission = 'd' if os.path.isdir(full_path) else 'f'
                        permission += 'r' if os.access(full_path, os.R_OK) else '-'
                        permission += 'w' if os.access(full_path, os.W_OK) else '-'
                        permission += 'x' if os.access(full_path, os.X_OK) else '-'
                        ans[file] = permission
                    return chr(0x14) + json.dumps(ans)

            elif ctrl == 0x15:
                try:
                    stat_result = os.stat(message)
                    stat_dict = {
                        "size": stat_result.st_size, "mtime": stat_result.st_mtime,
                        "ctime": stat_result.st_ctime, "mode": stat_result.st_mode,
                        "uid": stat_result.st_uid, "gid": stat_result.st_gid,
                    }
                    return chr(0x15) + json.dumps(stat_dict)
                except Exception as e:
                    return chr(0x15) + json.dumps({"error": str(e)})

            elif ctrl == 0x16:
                folder_path = message.strip()
                try:
                    if not os.path.exists(folder_path):
                        os.makedirs(folder_path)
                        return chr(0x16) + f"目录创建成功: {folder_path}"
                    else:
                        return chr(0x16) + f"目录已存在: {folder_path}"
                except Exception as e:
                    return chr(0x16) + f"目录创建失败: {e}"

            elif ctrl == 0x17:
                try:
                    old_path, new_path = message.split(' ', 1)
                    if not os.path.exists(old_path):
                        return chr(0x17) + f"源路径不存在: {old_path}"
                    elif os.path.exists(new_path):
                        return chr(0x17) + f"目标路径已存在: {new_path}"
                    else:
                        os.rename(old_path, new_path)
                        return chr(0x17) + f"重命名成功: {old_path} -> {new_path}"
                except Exception as e:
                    return chr(0x17) + f"重命名失败: {e}"
            else:
                return f"未知命令: {hex(ctrl)}\n"

        except Exception as e:
            return f"指令执行异常：{e}\n"

    def _send_file_chunks(self, session_id: str, file_path: str, chunk_count: int, client_id: str):
        try:
            with open(file_path, 'rb') as f:
                for chunk_index in range(chunk_count):
                    if session_id not in self.download_sessions:
                        print(
                            f" Download session {session_id} for client {client_id} cancelled during chunk sending.")
                        break
                    f.seek(chunk_index * CHUNK_SIZE)
                    chunk = f.read(CHUNK_SIZE)
                    if not chunk:
                        break
                    encoded_chunk = base64.b64encode(chunk).decode()
                    response_message = chr(0x18) + f"{session_id} {chunk_index} {encoded_chunk}"
                    static["event_handler"].trigger_event(
                        Event("ExternalIO_OUT", response_data=response_message, client_id=client_id)
                    )
            if session_id in self.download_sessions:
                static["event_handler"].trigger_event(
                    Event("ExternalIO_OUT", response_data=chr(0x18) + "END", client_id=client_id)
                )
                del self.download_sessions[session_id]
        except Exception as e:
            error_message = f"下载分块发送异常：{e}"
            static["event_handler"].trigger_event(
                Event("ExternalIO_ERROR", error_message=error_message, original_command="", client_id=client_id)
            )
            if session_id in self.download_sessions:
                del self.download_sessions[session_id]


external_handler_instance = None


def __init__():
    global external_handler_instance
    external_handler_instance = ExternalHandlerLogic()
    static["features"].update({
        "info_access": True,
        "file_transfer": True,
        "file_management": True,
        "terminal": True,
        "system_upgrade": True,
    })
